import torch
import torch.nn as nn


from model.net import Encoder, Decoder
from model.masknet import MaskNet


class SepFormer(nn.Module):
    def __init__(self, num_spks=2):
        super().__init__()
        self.encoder = Encoder()
        self.masknet = MaskNet()
        self.decoder = Decoder()
        self.num_spks = num_spks

    def forward(self, x):
        mix_w = self.encoder(x)
        est_mask = self.masknet(mix_w)
        mix_w = torch.stack([mix_w] * self.num_spks)
        sep_h = mix_w * est_mask
        if self.num_spks == 2:
            y1 = self.decoder(sep_h[0])
            y2 = self.decoder(sep_h[1])
        output = torch.stack([y1, y2], dim=1)
        return output
        pass

    pass


if __name__ == "__main__":
    x = torch.rand(1, 1, 32000)
    sepformer = SepFormer()
    with torch.no_grad():
        output = sepformer(x)
        print(output.shape)
    pass
